import os
from os.path import join, isdir
import sys
sys.path.append("..")
import argparse
from data_recording.multi_camera.parse_rosbag import parse_rosbag


def get_rosbag_fns(rosbag_dir):
    fns = []
    for fn in os.listdir(rosbag_dir):
        if not ".bag" in fn:
            continue
        fns.append(fn)
    fns.sort()
    return fns


def get_time(fn):
    """
    parse a fake time from rosbag filename ("YY-MM-DD-HH-MM-SS.bag")
    """
    fn = fn.split(".")[0]
    year, month, day, hour, minute, second = [int(x) for x in fn.split("-")]
    T = year * 100000000 + month * 5000000 + day * 86400 + hour * 3600 + minute * 60 + second
    return T


def align_rosbag_fns(fns1, fns2, threshould=5):  # 5s
    print(len(fns1), len(fns2))
    if (len(fns1) == 0) or (len(fns2) == 0):
        return []
    fns1.sort()
    fns2.sort()
    
    paired_fns = []
    N = len(fns2)
    p = 0
    for fn in fns1:
        t = get_time(fn)
        while (p + 1 < N) and (abs(t - get_time(fns2[p+1])) <= abs(t - get_time(fns2[p]))):
            p += 1
        if abs(t - get_time(fns2[p])) < threshould:
            paired_fns.append([fn, fns2[p]])
        
    return paired_fns


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--slot_name', type=str, default="20230804_1")
    args = parser.parse_args()
    return args


if __name__  == "__main__":
    
    #######################################################################
    #######################################################################
    
    args = parse_args()
    slot_name = args.slot_name
    
    dataset_dir = "/share/datasets/HHO_dataset"
    rosbag1_dir = join(dataset_dir, "rawdata", slot_name, "rosbag1")
    rosbag2_dir = join(dataset_dir, "rawdata", slot_name, "rosbag2")
    save_dir = join(dataset_dir, "data", slot_name)
    if isdir(save_dir):  # move
        os.makedirs(join(dataset_dir, "data_old"), exist_ok=True)
        os.system("mv {} {}".format(save_dir, join(dataset_dir, "data_old", slot_name)))
    os.makedirs(save_dir, exist_ok=True)
    
    rosbag1_fns = get_rosbag_fns(rosbag1_dir)
    rosbag2_fns = get_rosbag_fns(rosbag2_dir)
    paired_fns = align_rosbag_fns(rosbag1_fns, rosbag2_fns, threshould=5)
    
    with open(join(save_dir, "aligned_rosbag_filenames.txt"), "w") as f:
        for x in paired_fns:
            f.write(x[0] + "," + x[1] + "\n")
    
    for (i, paired_fn) in enumerate(paired_fns):
        rosbag1_path = join(rosbag1_dir, paired_fn[0])
        rosbag2_path = join(rosbag2_dir, paired_fn[1])
        clip_save_dir = join(save_dir, str(i).zfill(3))
        
        # parse data
        parse_rosbag(rosbag2_path, clip_save_dir, save_video=True)
        parse_rosbag(rosbag1_path, clip_save_dir, save_video=True)
